package com.mydemo.utils.base;

import org.vosk.Model;
import org.vosk.Recognizer;

import javax.sound.sampled.AudioFormat;
import javax.sound.sampled.AudioInputStream;
import javax.sound.sampled.AudioSystem;
import javax.sound.sampled.UnsupportedAudioFileException;
import java.io.*;
import java.nio.charset.StandardCharsets;

public class VoskAudioTranscriber {

    public static void main(String[] args) {
        // 1. 模型路径（请替换为你的实际路径）
        // String modelPath = "path/to/your/vosk-model-cn-kaldi-multicn-0.15";
        // String modelPath = "D:\\mp4totext\\vosk-model-cn-kaldi-multicn-0.15";
        String modelPath = "D:\\mp4totext\\vosk-model-cn-0.22";
        // 2. 待识别的音频文件路径
        // String audioFilePath = "path/to/your/audio.wav";
        String audioFilePath = "D:\\mp4totext\\古籍修复.wav";
        // 3. 输出文本文件路径
        // String outputTextPath = "path/to/your/transcription_result.txt";
        String outputTextPath ="D:\\mp4totext\\古籍修复-022.txt";

        Model model = null;
        AudioInputStream audioInputStream = null;
        BufferedWriter writer = null;

        try {
            // 加载模型
            System.out.println("正在加载语音识别模型...");
            model = new Model(modelPath);
            Recognizer recognizer = new Recognizer(model, 16000.0f);

            // 检查音频文件是否存在
            File audioFile = new File(audioFilePath);
            if (!audioFile.exists()) {
                System.err.println("错误：音频文件未找到！路径：" + audioFilePath);
                return;
            }

            // 创建输出文本文件
            writer = new BufferedWriter(new OutputStreamWriter(
                    new FileOutputStream(outputTextPath), StandardCharsets.UTF_8));

            // 配置并获取音频输入流
            audioInputStream = AudioSystem.getAudioInputStream(audioFile);
            AudioFormat targetFormat = new AudioFormat(16000, 16, 1, true, false);
            AudioInputStream convertedAudio = AudioSystem.getAudioInputStream(targetFormat, audioInputStream);

            // 读取音频数据并识别
            int bytesRead;
            byte[] buffer = new byte[4096];
            StringBuilder fullText = new StringBuilder();

            System.out.println("开始语音识别...");

            while ((bytesRead = convertedAudio.read(buffer)) >= 0) {
                if (bytesRead > 0) {
                    if (recognizer.acceptWaveForm(buffer, bytesRead)) {
                        // 获取完整句子结果
                        String result = recognizer.getResult();
                        String text = extractTextFromResult(result);

                        if (!text.isEmpty()) {
                            // 在控制台打印
                            System.out.println("识别结果: " + text);

                            // 写入文件
                            writer.write(text);
                            writer.newLine();
                            writer.flush();

                            // 添加到完整文本
                            fullText.append(text).append(" ");
                        }
                    } else {
                        // 可选：打印部分识别结果用于调试
                        // String partial = recognizer.getPartialResult();
                        // System.out.println("部分结果: " + partial);
                    }
                }
            }

            // 处理最终的识别结果
            String finalResult = recognizer.getFinalResult();
            String finalText = extractTextFromResult(finalResult);

            if (!finalText.isEmpty()) {
                System.out.println("最终识别结果: " + finalText);
                writer.write(finalText);
                writer.newLine();
                fullText.append(finalText);
            }

            // 打印完整转录文本
            String completeTranscription = fullText.toString().trim();
            if (!completeTranscription.isEmpty()) {
                System.out.println("\n===== 完整转录内容 =====");
                System.out.println(completeTranscription);
                System.out.println("========================");

                // 将完整内容写入文件开头
                writer.write("===== 完整转录内容 =====\n");
                writer.write(completeTranscription);
                writer.write("\n========================\n");
            }

            System.out.println("\n语音识别完成！结果已保存至: " + outputTextPath);

        } catch (UnsupportedAudioFileException e) {
            System.err.println("错误：不支持的音频文件格式！");
            System.err.println("请确保音频文件为WAV格式，并使用FFmpeg转换为16kHz, 16bit, 单声道格式");
            e.printStackTrace();
        } catch (IOException e) {
            System.err.println("错误：文件读写异常！");
            e.printStackTrace();
        } catch (Exception e) {
            System.err.println("错误：语音识别过程中发生异常！");
            e.printStackTrace();
        } finally {
            // 关闭资源
            try {
                if (audioInputStream != null) {
                    audioInputStream.close();
                }
                if (writer != null) {
                    writer.close();
                }
                if (model != null) {
                    model.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 从Vosk识别结果JSON中提取纯文本内容
     * @param resultJson Vosk返回的JSON结果字符串
     * @return 提取出的纯文本内容
     */
    private static String extractTextFromResult(String resultJson) {
        try {
            // 简单的JSON解析，提取"text"字段的值
            // 结果格式通常为: {"text": "识别的文字内容"}
            int textStart = resultJson.indexOf("\"text\"") + 8;
            int textEnd = resultJson.lastIndexOf("\"");
            if (textStart > 7 && textEnd > textStart) {
                return resultJson.substring(textStart, textEnd);
            }
        } catch (Exception e) {
            System.err.println("解析识别结果时出错: " + resultJson);
        }
        return "";
    }
}
